import numpy as np
import nibabel as nib
from scipy import ndimage
import os
import glob
import traceback

def normalize_mri_dimensions(image_path, mask_path, output_image_dir, output_mask_dir, target_shape=(256, 256, 256)):
    """
    统一MRI图像和mask的xyz三个维度
    
    参数:
    - image_path: MRI图像文件路径
    - mask_path: mask文件路径  
    - output_image_dir: 输出图像目录
    - output_mask_dir: 输出mask目录
    - target_shape: 目标形状 (x, y, z)
    """
    
    # 创建输出目录
    os.makedirs(output_image_dir, exist_ok=True)
    os.makedirs(output_mask_dir, exist_ok=True)
    
    # 加载图像和mask
    try:
        image_nii = nib.load(image_path)
        mask_nii = nib.load(mask_path)
    except Exception as e:
        print(f"加载文件时出错: {str(e)}")
        return None, None
    
    image_data = image_nii.get_fdata()
    mask_data = mask_nii.get_fdata()
    
    print(f"原始图像形状: {image_data.shape}")
    print(f"原始mask形状: {mask_data.shape}")
    
    # 检查维度是否匹配
    if image_data.shape != mask_data.shape:
        print(f"警告: 图像和mask形状不匹配! 图像: {image_data.shape}, mask: {mask_data.shape}")
        
        # 如果维度数相同但大小不同，调整mask大小
        if len(image_data.shape) == len(mask_data.shape):
            try:
                mask_data = resize_to_target(mask_data, image_data.shape, is_mask=True)
            except Exception as e:
                print(f"调整mask大小时出错: {str(e)}")
                return None, None
    
    # 处理长宽比过大的问题
    try:
        processed_image = process_large_aspect_ratio(image_data, target_shape)
        processed_mask = process_large_aspect_ratio(mask_data, target_shape, is_mask=True)
    except Exception as e:
        print(f"处理长宽比时出错: {str(e)}")
        return None, None
    
    print(f"处理后图像形状: {processed_image.shape}")
    print(f"处理后mask形状: {processed_mask.shape}")
    
    # 获取基础文件名
    image_base_name = os.path.basename(image_path)
    mask_base_name = os.path.basename(mask_path)
    
    # 创建新的NIfTI图像
    try:
        processed_image_nii = nib.Nifti1Image(processed_image, image_nii.affine)
        processed_mask_nii = nib.Nifti1Image(processed_mask, mask_nii.affine)
    except Exception as e:
        print(f"创建NIfTI图像时出错: {str(e)}")
        return None, None
    
    # 保存文件到各自的目录
    output_image_path = os.path.join(output_image_dir, image_base_name)
    output_mask_path = os.path.join(output_mask_dir, mask_base_name)
    
    try:
        nib.save(processed_image_nii, output_image_path)
        nib.save(processed_mask_nii, output_mask_path)
    except Exception as e:
        print(f"保存文件时出错: {str(e)}")
        return None, None
    
    print(f"图像保存到: {output_image_path}")
    print(f"mask保存到: {output_mask_path}")
    return output_image_path, output_mask_path

def process_large_aspect_ratio(data, target_shape, is_mask=False):
    """
    处理长宽比过大的数据
    
    参数:
    - data: 输入数据
    - target_shape: 目标形状
    - is_mask: 是否是mask数据
    """
    
    current_shape = data.shape
    aspect_ratio = max(current_shape) / min(current_shape)
    print(f"当前形状: {current_shape}, 长宽比: {aspect_ratio:.2f}")
    
    # 检查是否需要处理长宽比
    if aspect_ratio > 10:  # 长宽比阈值
        print(f"检测到长宽比过大: {aspect_ratio:.2f}, 进行标准化处理")
    
    # 调整到目标形状
    if current_shape != target_shape:
        return resize_to_target(data, target_shape, is_mask=is_mask)
    
    return data

def resize_to_target(data, target_shape, is_mask=False):
    """
    将数据调整到目标形状
    
    参数:
    - data: 输入数据
    - target_shape: 目标形状
    - is_mask: 是否是mask数据（决定使用哪种插值方法）
    """
    
    # 计算缩放因子
    zoom_factors = [target_shape[i] / data.shape[i] for i in range(len(data.shape))]
    
    try:
        if is_mask:
            # 对于mask使用最近邻插值，保持标签完整性
            resized_data = ndimage.zoom(data, zoom_factors, order=0, mode='nearest')
        else:
            # 对于图像使用3阶插值
            resized_data = ndimage.zoom(data, zoom_factors, order=3, mode='constant', cval=0)
    except Exception as e:
        print(f"缩放数据时出错: {str(e)}")
        raise
    
    return resized_data

def extract_base_id(filename):
    """
    从文件名中提取基础ID
    
    参数:
    - filename: 文件名
    
    返回:
    - 基础ID (如 "877033_T2_axi")
    """
    # 移除文件扩展名
    base_name = os.path.splitext(filename)[0]
    if base_name.endswith('.nii'):
        base_name = os.path.splitext(base_name)[0]
    
    # 移除后缀部分 (_000 或 _roi)
    if base_name.endswith('_000'):
        base_name = base_name[:-4]
    elif base_name.endswith('_roi'):
        base_name = base_name[:-4]
    
    return base_name

def find_corresponding_mask(image_filename, mask_dir):
    """
    在mask目录中查找对应的mask文件
    
    参数:
    - image_filename: 图像文件名
    - mask_dir: mask目录路径
    """
    
    # 提取基础ID
    base_id = extract_base_id(image_filename)
    
    # 可能的mask文件名模式
    possible_names = [
        f"{base_id}_roi.nii.gz",
        f"{base_id}_roi.nii",
        f"{base_id}_mask.nii.gz",
        f"{base_id}_mask.nii",
        f"{base_id}_seg.nii.gz",
        f"{base_id}_seg.nii",
    ]
    
    for mask_name in possible_names:
        mask_path = os.path.join(mask_dir, mask_name)
        if os.path.exists(mask_path):
            return mask_path
    
    # 如果没有找到精确匹配，尝试模糊匹配
    mask_files = glob.glob(os.path.join(mask_dir, '*'))
    for mask_file in mask_files:
        mask_filename = os.path.basename(mask_file)
        mask_base_id = extract_base_id(mask_filename)
        
        # 检查是否匹配基础ID
        if base_id == mask_base_id:
            return mask_file
    
    return None

def batch_process_mri_separate_folders(image_dir, mask_dir, output_image_dir, output_mask_dir, target_shape=(256, 256, 256)):
    """
    批量处理MRI图像，分别从图像和mask文件夹读取
    
    参数:
    - image_dir: 原图像目录
    - mask_dir: mask目录
    - output_image_dir: 输出图像目录
    - output_mask_dir: 输出mask目录
    - target_shape: 目标形状
    """
    
    # 支持的图像格式
    image_extensions = ['*.nii.gz', '*.nii']
    
    # 查找所有图像文件
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(image_dir, ext)))
    
    print(f"找到 {len(image_files)} 个图像文件")
    
    processed_files = []
    failed_files = []
    
    # 添加进度计数器
    total_files = len(image_files)
    
    for i, image_file in enumerate(image_files):
        # 获取文件名（不含路径）
        image_filename = os.path.basename(image_file)
        
        # 显示进度
        print(f"\n进度: {i+1}/{total_files} ({((i+1)/total_files)*100:.1f}%)")
        
        # 查找对应的mask文件
        mask_file = find_corresponding_mask(image_filename, mask_dir)
        
        if mask_file and os.path.exists(mask_file):
            print(f"处理文件: {image_filename}")
            print(f"匹配mask: {os.path.basename(mask_file)}")
            try:
                output_image, output_mask = normalize_mri_dimensions(
                    image_file, mask_file, output_image_dir, output_mask_dir, target_shape
                )
                if output_image and output_mask:
                    processed_files.append((output_image, output_mask))
                    print(f"✓ 成功处理: {image_filename}")
                else:
                    failed_files.append((image_file, mask_file))
                    print(f"✗ 处理失败: {image_filename}")
            except Exception as e:
                failed_files.append((image_file, mask_file))
                print(f"✗ 处理文件 {image_filename} 时出错: {str(e)}")
                traceback.print_exc()
        else:
            failed_files.append((image_file, None))
            print(f"✗ 未找到 {image_filename} 对应的mask文件")
    
    print(f"\n批量处理完成!")
    print(f"成功处理: {len(processed_files)} 对文件")
    print(f"处理失败: {len(failed_files)} 对文件")
    
    # 如果有失败的文件，保存失败列表
    if failed_files:
        failed_log_path = os.path.join(output_image_dir, "failed_files.log")
        with open(failed_log_path, 'w') as f:
            f.write("失败文件列表:\n")
            for image_file, mask_file in failed_files:
                f.write(f"图像: {image_file}, Mask: {mask_file}\n")
        print(f"失败文件列表已保存到: {failed_log_path}")
    
    return processed_files, failed_files

def validate_file_pairs(image_dir, mask_dir):
    """
    验证图像和mask文件的配对情况
    
    参数:
    - image_dir: 图像目录
    - mask_dir: mask目录
    """
    
    image_extensions = ['*.nii.gz', '*.nii']
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(image_dir, ext)))
    
    print("文件配对验证:")
    print("=" * 50)
    
    matched_pairs = []
    unmatched_images = []
    
    for image_file in image_files:
        image_filename = os.path.basename(image_file)
        mask_file = find_corresponding_mask(image_filename, mask_dir)
        
        if mask_file:
            matched_pairs.append((image_file, mask_file))
            print(f"✅ 匹配: {image_filename} -> {os.path.basename(mask_file)}")
        else:
            unmatched_images.append(image_file)
            print(f"❌ 未匹配: {image_filename}")
    
    print(f"\n统计:")
    print(f"总图像文件: {len(image_files)}")
    print(f"成功匹配: {len(matched_pairs)}")
    print(f"未匹配: {len(unmatched_images)}")
    
    return matched_pairs, unmatched_images

def analyze_aspect_ratios(data_dir, file_pattern="*.nii.gz"):
    """
    分析目录中所有文件的长宽比
    """
    
    files = glob.glob(os.path.join(data_dir, file_pattern))
    
    aspect_ratios = []
    
    for file_path in files:
        try:
            nii = nib.load(file_path)
            data = nii.get_fdata()
            shape = data.shape
            
            # 计算最大长宽比
            aspect_ratio = max(shape) / min(shape)
            aspect_ratios.append({
                'file': os.path.basename(file_path),
                'shape': shape,
                'aspect_ratio': aspect_ratio
            })
            
            if aspect_ratio > 10:
                print(f"⚠️  高长宽比: {file_path} - 形状: {shape} - 比例: {aspect_ratio:.2f}")
                
        except Exception as e:
            print(f"分析文件 {file_path} 时出错: {str(e)}")
    
    return aspect_ratios

# 使用示例
if __name__ == "__main__":
    # 设置路径
    image_dir =r"i:\train\img"  # 原图像文件夹
    mask_dir = r"i:\train\roi"    # mask文件夹
    output_image_dir = r"i:\train\dimensions_img"  # 处理后的图像输出文件夹
    output_mask_dir = r"i:\train\dimensions_roi"    # 处理后的mask输出文件夹
    
    # 首先验证文件配对
    print("开始验证文件配对...")
    matched_pairs, unmatched_images = validate_file_pairs(image_dir, mask_dir)
    
    if unmatched_images:
        print(f"\n警告: 有 {len(unmatched_images)} 个图像文件没有找到对应的mask")
        proceed = input("是否继续处理? (y/n): ")
        if proceed.lower() != 'y':
            exit()
    
    # 批量处理
    print("\n开始批量处理...")
    processed_files, failed_files = batch_process_mri_separate_folders(
        image_dir=image_dir,
        mask_dir=mask_dir,
        output_image_dir=output_image_dir,
        output_mask_dir=output_mask_dir,
        target_shape=(256, 256, 256)  # 可根据需要调整
    )
    
    print(f"\n处理完成!")
    print(f"成功处理: {len(processed_files)} 对文件")
    print(f"处理失败: {len(failed_files)} 对文件")
    print(f"处理后的图像保存在: {output_image_dir}")
    print(f"处理后的mask保存在: {output_mask_dir}")
    
    # 分析长宽比
    print("\n分析图像长宽比...")
    analyze_aspect_ratios(output_image_dir)
    print("\n分析mask长宽比...")
    analyze_aspect_ratios(output_mask_dir)

